#-------------------------------------------------------------------------------
# Name:        detect
# Purpose:
#
# Author:      Administrator
#
# Created:     15/04/2024
# Copyright:   (c) Administrator 2024
# Licence:     <your licence>
#-------------------------------------------------------------------------------
import argparse
import csv
import os
import platform
import sys
from pathlib import Path
from ultralytics.utils.plotting import Annotator, colors, save_one_box
import torch
import os

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
from utils.general import (
    LOGGER,
    Profile,
    check_file,
    check_img_size,
    check_imshow,
    check_requirements,
    colorstr,
    cv2,
    increment_path,
    non_max_suppression,
    print_args,
    scale_boxes,
    strip_optimizer,
    xyxy2xywh,
)
from utils.torch_utils import select_device, smart_inference_mode
from models.common import DetectMultiBackend
import np
import numpy
from PIL import Image, ImageDraw, ImageFont
import time
from utils.augmentations import letterbox

class SingletonMeta(type):
    _instances = {}

    def __call__(cls, *args, **kwargs):
        if cls not in cls._instances:
            cls._instances[cls] = super().__call__(*args, **kwargs)
        return cls._instances[cls]

class Singleton(metaclass=SingletonMeta):
    _classes = {}

    def __init__(self, class_name, class_instance):
        self.class_name = class_name
        self.class_instance = class_instance

    @classmethod
    def register_class(cls, class_name, class_instance):
        cls._classes[class_name] = cls(class_name, class_instance)

    @classmethod
    def get_class_instance(cls, class_name):
        return cls._classes.get(class_name, None)

@smart_inference_mode()
class MyModel():
    Model=None
    opt=None
    def __init__(self,weights=''):
        start_time = time.time()
        parser = argparse.ArgumentParser()
        parser.add_argument("--weights", nargs="+", type=str, default="", help="model path or triton URL")
        parser.add_argument('--source', type=str, default="D:/aaa/", help='source')  # file/folder, 0 for webcam
        parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
        parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
        parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
        parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
        parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
        parser.add_argument("--view-img", action="store_true", help="show results")
        parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
        parser.add_argument("--save-csv", action="store_true", help="save results in CSV format")
        parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
        parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
        parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
        parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
        parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
        parser.add_argument("--augment", action="store_true", help="augmented inference")
        parser.add_argument("--visualize", action="store_true", help="visualize features")
        parser.add_argument("--update", action="store_true", help="update all models")
        parser.add_argument("--project", default=ROOT / "runs/detect", help="save results to project/name")
        parser.add_argument("--name", default="exp", help="save results to project/name")
        parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
        parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
        parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
        parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
        parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
        parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
        parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
        self.opt = parser.parse_args()
        self.opt.imgsz *= 2 if len(self.opt.imgsz) == 1 else 1  # expand

        print(self.opt)

        device = select_device("")
        dnn=False
        half=False
        try:
            self.Model = DetectMultiBackend(weights, device=device, dnn=dnn, data="", fp16=half)
        except Exception as e:
            print(e)

        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"model时间：{elapsed_time} 秒 ")

    def model(self):
        return self.Model

    def parser(self):
        return self.opt


def run(
    weights=ROOT / "",  # model path or triton URL
    source=ROOT / "../test",  # file/dir/URL/glob/screen/0(webcam)
    data=ROOT / "data/coco128.yaml",  # dataset.yaml path
    imgsz=(640, 640),  # inference size (height, width)
    conf_thres=0.25,  # confidence threshold
    iou_thres=0.45,  # NMS IOU threshold
    max_det=1000,  # maximum detections per image
    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
    view_img=False,  # show results
    save_txt=False,  # save results to *.txt
    save_csv=False,  # save results in CSV format
    save_conf=False,  # save confidences in --save-txt labels
    save_crop=False,  # save cropped prediction boxes
    nosave=False,  # do not save images/videos
    classes=None,  # filter by class: --class 0, or --class 0 2 3
    agnostic_nms=False,  # class-agnostic NMS
    augment=False,  # augmented inference
    visualize=False,  # visualize features
    update=False,  # update all models
    project=ROOT / "runs/detect",  # save results to project/name
    name="exp",  # save results to project/name
    exist_ok=False,  # existing project/name ok, do not increment
    line_thickness=3,  # bounding box thickness (pixels)
    hide_labels=False,  # hide labels
    hide_conf=False,  # hide confidences
    half=False,  # use FP16 half-precision inference
    dnn=False,  # use OpenCV DNN for ONNX inference
    vid_stride=1,  # video frame-rate stride
    model=None
):
    start_time = time.time()
    source = str(source)
    save_img = not nosave and not source.endswith(".txt")  # save inference images
    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
    is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://"))
    webcam = source.isnumeric() or source.endswith(".streams") or (is_url and not is_file)
    screenshot = source.lower().startswith("screen")
    if is_url and is_file:
        source = check_file(source)  # download

    my_model_instance = Singleton.get_class_instance('MyModel')
    model=my_model_instance.class_instance.model()

    stride, names, pt = model.stride, model.names, model.pt
    imgsz = check_img_size(imgsz, s=stride)  # check image size

    # Dataloader
    bs = 1  # batch_size
    if webcam:
        view_img = check_imshow(warn=True)
        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
        bs = len(dataset)
    elif screenshot:
        dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
    else:
        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference
    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
    seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))

    print('dt===',dt)

    j=0

    for path, im, im0s, vid_cap, s in dataset:

        with dt[0]:
            im = torch.from_numpy(im).to(model.device)
            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim
            if model.xml and im.shape[0] > 1:
                ims = torch.chunk(im, im.shape[0], 0)


        pred = model(im, augment=augment, visualize=visualize)
        # NMS
        with dt[2]:
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)



        for i, det in enumerate(pred):
            p, s, im0 = source, '', im0s

            # print('im0_shape：', im0.shape)

            s += '%gx%g ' % im.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  #  normalization gain whwh

            if det is not None and len(det):
                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, 5].unique():
                    n = (det[:, 5] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                output_dict_ = []
                for *xyxy, conf, cls in det:
                    x1, y1, x2, y2 = xyxy
                    print(conf,cls)
                    output_dict_.append((float(x1), float(y1), float(x2), float(y2)))
                    label = '%s %.2f' % (names[int(cls)], conf)
                    print(j,"---------------------------------------------------------------------")
                    print(path)
                    print("尺寸：", im0.shape)
                    print("坐标：", (float(x1), float(y1), float(x2), float(y2)))
                    print("标识：", label)

            plot_square(det,im0,names)

        j+=1


def cv2ImgAddText(img, text, left, top, textColor=(0, 255, 0), textSize=20):

    if (isinstance(img, numpy.ndarray)):
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))

    draw = ImageDraw.Draw(img)
    fontText = ImageFont.truetype(
        "font/simsun.ttc", textSize, encoding="utf-8")
    draw.text((left, top), text, textColor, font=fontText)

    return cv2.cvtColor(numpy.asarray(img), cv2.COLOR_RGB2BGR)

def plot_square(det,image_pil,names):

    image_np = np.array(image_pil)

    image = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)

    for *xyxy, conf, cls in det:
        x1, y1, x2, y2 = xyxy
        print(conf,cls)
        label = '%s %.2f' % (names[int(cls)], conf)
        top_left = int(x1), int(y1)
        bottom_right = int(x2), int(y2)

        cv2.rectangle(image_pil, top_left, bottom_right, (0, 255, 0), 3)

        image_pil = cv2ImgAddText(image_pil, label, int(x1), int(y1) - 30, (0, 255, 0), 30)

    cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
    cv2.imshow('Image', image_pil)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def detect_image(im0, model, device='cpu', conf_thres=0.25, iou_thres=0.45,show=0):

    if(model is None):
        return 0,0,None

    imgsz=(640, 640)
    stride, names, pt = model.stride, model.names, model.pt
    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))  # warmup
    seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))

    augment=False
    visualize=False
    classes=None
    agnostic_nms=False
    max_det=1000

    im = letterbox(im0, 640, 32, True)[0]  # padded resize
    im = im.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    im = np.ascontiguousarray(im)  # contiguous

    im0s=im0
    vid_cap=None
    s=''
    with dt[0]:
            im = torch.from_numpy(im).to(model.device)
            im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32
            im /= 255  # 0 - 255 to 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim
            if model.xml and im.shape[0] > 1:
                ims = torch.chunk(im, im.shape[0], 0)


    pred = model(im, augment=augment, visualize=visualize)
        # NMS
    with dt[2]:
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)

    for i, det in enumerate(pred):
            s, im0 =  '', im0s

            # print('im0_shape：', im0.shape)

            s += '%gx%g ' % im.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  #  normalization gain whwh

            if det is not None and len(det):
                det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results
                for c in det[:, 5].unique():
                    n = (det[:, 5] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results
                output_dict_ = []


                if(show==1):
                    plot_square(det,im0,names)

    return names,im0,det

def st(directory,weights=None):
    start_time = time.time()
    current_directory = os.getcwd()
    Singleton.register_class('MyModel', MyModel(weights))

    my_model_instance = Singleton.get_class_instance('MyModel')
    model=my_model_instance.class_instance.model()
    opt=my_model_instance.class_instance.parser()

    #run(**vars(opt))

    files = os.listdir(directory)

    for file in files:
        image_path = directory+file
        print('image_path=',image_path)
        im0 = cv2.imread(image_path)
        names,im0,det=detect_image(im0, model, device='cpu',show=1)

        if(det is None):
            continue
        for *xyxy, conf, cls in det:
            x1, y1, x2, y2 = xyxy
            label = '%s' % (names[int(cls)])
            probability='%.2f' % (conf)
            top_left = int(x1), int(y1)
            bottom_right = int(x2), int(y2)
            print(label,probability,top_left,bottom_right)

    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"all时间：{elapsed_time} 秒 ")
